########################################################################
##
## functions for doing some simple time-series related manipulations,
## including an iterative Cochrane-Orcutt correction for AR(1)
## disturbances in the linear regression model
##
## the functions are:
##
## mylag(x,lag=1) -- create new vars, lag of input var, default is lag 1
##
## mydw(lmobj)  -- Durbin-Watson statistic for an object created by lm()
##                 and p-value, if lmtest library is available
##
## myrho(lmobj) -- Estimate of 1st order serial correlation
##                 in resids of lm() object
##
## mystar(x,rho) -- given data x and rho, create transformed data by subtracting
##                  rho times lagged value of x; does the Prais-Winsten transform
##                  for the first observation
##
## ar1plot(lmobj) -- plot resids against lagged resids, given an object from lm()
##
## cochrane.orcutt(lmobj) -- Cochrane-Orcutt iterative EGLS, correcting for
##                           1st order serial correlation
##
## durbinh(lmobj,"lagvar") -- Durbin's h, testing for 1st order serial correlation
##                            among residuals of lm() object, where the lm object
##                            has "lagvar" as a regressor (a variable containing
##                            the lag of the dependent variable).
##
## impulse(beta=1,rho) -- plot decay of a time zero shock of size beta,
##                        rho is AR(1) decay
##
## to use these functions, simply source the file into your R session
## > source("ar1.r")
##
## Simon Jackman, Dept of Political Science, Stanford University
## 1999, 2000, 2002, 2004.
#########################################################################
## a function for creating lags properly
## usage: xlag <- mylag(x)
## where x can be either a vector or a matrix
#########################################################################
mylag <- function(x,lag=1){
  # is x a matrix?
  if (is.matrix(x)){
    n <- dim(x)[1]
    k <- dim(x)[2]
    if(k>1){
      z <- rep(NA,k)         # 1st observations are missing
      y <- x[1:(n-lag),]    
      for (i in 1:lag)
        y <- rbind(z,y)
    }
  }
  else # assume x is a vector
    {
      n <- length(x)
      y <- c(rep(NA,lag),x[1:(n-lag)])
    }
  y
}

############################################################################
## a function for detecting AR(1) errors
## to use: run a regression with the lm command then feed the lm.object
## to this function
## e.g.,
## blah <- lm(y~x,data=junk)
## myrho(blah)
###########################################################################
myrho <- function(fit,print=F){
  e <- resid(fit)
  e1 <- mylag(e)
  n <- length(e)
  rho <- sum(e[-1]*e[-n])/sum(e^2)
  #dwdat <- as.data.frame(list(e=e,e1=e1))
  #dwreg <- lm(e~-1+e1,data=dwdat,na.action=na.omit)
  #if(print)
  #print(summary(dwreg))
  #rho <- coef(dwreg)
  rho
}

############################################################################
# a function for transforming regression variables for GLS, given rho
#
# usage: ystar <- mystar(y,rho)
###########################################################################
mystar <- function(x,rho){
    xstar <- x-(rho*mylag(x,1))
    ## is x a matrix?
    if (is.matrix(x))
      xstar[1,] <- sqrt(1-(rho^2))*x[1,]         # PW transform
    else
      xstar[1] <- sqrt(1-(rho^2))*x[1]
    xstar
}

############################################################################
## ar1plot
## 
## diagnostic plot, check for AR(1) pattern in the residuals of
## lm object, by plotting \hat{u}_t against \hat{u}_{t-1} and overlay
## estimate of rho
##
## simon jackman, dept of political science, stanford university
## march 2000
############################################################################
ar1plot <- function(fit){
  e <- resid(fit)
  e1 <- mylag(resid(fit))
  plot(x=e1,y=e,
       xlab="Lagged Residual",
       ylab="Residual")
  rho <- myrho(fit)
  abline(a=0,b=rho)
  cat(paste("rho =",rho,"\n"))
  invisible(NULL)
}
     
##################################################################
## cochrane.orcutt
##
## a function to do iterative Cochran-Orcutt for AR(1) errors
## usage:
##   > lmobject <- lm(y~x)
##   > blah <- cochrane.orcutt(lmobject)
##
## optional arguments:
## itermax -- maximum number of Cochrane-Orcutt iterations
##            (default is 20)
## eps     -- covergence criterion (default is .000001);
##            iterations cease if rho
##            changes by less than this amount in
##            successive iterations
##
## simon jackman, dept of political science,
## stanford university
## winter 1999, amended march 2000
#################################################################
cochrane.orcutt <- function(fit, itermax = 20, eps = 1e-06)
{
  cat("     Iterative Cochrane-Orcutt EGLS\n")
  cat("fitting regression with AR(1) disturbances\n")
  cat("------------------------------------------\n\n")
  cat(paste("will iterate for max of",itermax,"iterations\n"))
  cat(paste("or until AR(1) parameter changes by less than",eps,"\n"))
  x <- model.matrix(fit)
  xnames <- dimnames(x)[[2]]
  y <- fitted(fit) + resid(fit)
  n <- length(y)
  k <- dim(x)[2]
  rho <- myrho(fit)
                                        # starting estimate of rho
  tss <- crossprod(y - mean(y))         # total sum of squares
  iota <- rep(1,n)
  M <- diag(n) - 1/n * iota%*%t(iota)   # centering matrix

  rhostar <- rho                        # initialize rhostar
  for(iter in 1:itermax) {
    ystar <- mystar(y, rhostar)         # transform y, including PW
    xstar <- mystar(x, rhostar)         # transform x, including PW
    localdata <- as.data.frame(list(ystar = ystar, xstar = xstar))
    regobj <- lm(ystar~-1+xstar,
                 data=localdata)
    bstar <- coef(regobj)
    e <- resid(regobj)                  # white noise residuals
    epe <- crossprod(e)                 # sum sq white noise residuals

    yhat <- x %*% bstar                 # predictions, using original data
    u <- y - yhat                       # regression residuals

    ## r-squared calculations, yuk, trying to replicate
    ## the SAS Reg R-Squared stuff p 183 of ETS manual
    tsst <- ystar[-1] - xstar[-1,1]*bstar[1]
    tsst <- crossprod(scale(tsst,scale=F))
    tssp <- apply(xstar[-1,],2,function(x)x-mean(x))
    tssp <- tssp[,-1]
    tssp <- tssp%*%matrix(bstar[-1],ncol=1)
    tssp <- crossprod(scale(tssp,scale=F))
    r2 <- tssp/tsst

    ## alternative rsquared calcs
    ssy <- ystar - bstar[1]*xstar[,1]
    
    sigma2 <- epe/(n - k)               # white noise variance
                                        # log-likelihood, Greene, p600
    llh <- (-n/2)*(log(2*pi)+log(sigma2)) +
      (1/2)*(log(1 - (rhostar^2))) -
        crossprod(e/(2*sigma2))

    cat(paste("Iteration ", iter,
              ": rho = ",    signif(rho, 3),
              " rhostar = ", signif(rhostar, 3),
              " llh = ",     signif(llh, 4),
              " ssu = ",     signif(epe, 4),
              "\n", sep = ""))
    rho <- sum(u[-1] * u[-n])/crossprod(u[-n])  # Gujarati, 3rd ed, p431
    rho <- as.numeric(rho)
    if ((rho-rhostar)<eps)              # check for convergence
      break
    else
      rhostar <- rho
  }
  rhose <- sqrt(sigma2/crossprod(u[-n]))      # std error on rho
  out <- list(fit = regobj,
              rho = rhostar,
              rhose = rhose,
              r2 = r2,
              ystar = ystar,
              xstar = xstar,
              ssy = ssy)

  cat("Iterations Terminating...\n")
  cat("\n")
  cat("EGLS Results:\n")
  
  names(out$fit$coefficients) <- xnames
  print(summary(out$fit)$coefficients)
  cat(paste("Estimate of rho:",
            signif(out$rho,4),"(disturbance AR1 parameter)\n"))
  cat(paste("StdError of rho:",
            signif(out$rhose,4),"\n"))
  cat(paste("regression r-squared after AR(1) correction:",
            signif(out$r2,4),"\n"))
  invisible(NULL)          
  out
}

############################################################################
## a function for calculating Durbin's h test
## to use: run a regression with a lagged y variable to the lm command then
## feed the lm.object and the name of the lagged variable to this function
## e.g.,
## blah <- lm(y ~ mylag + x,data=junk)
## durbinh(blah, "mylag")  ## n.b., a string indicating name of lagvar
## Christina Maimone (edited by Simon Jackman)
## Dept of Political Science, Stanford Univ, March 1, 2002
#############################################################################

durbinh <- function(ols, lagvar) {
  p <- myrho(ols)
  t <- length(resid(ols))
  v <- ((summary(ols))$cov)[lagvar,lagvar]
  h <- p*((t/(1-t*v))^.5)
  cat(paste("Disturbance AR(1) parameter:",
            round(p,2),"\n"))
  cat(paste("Durbin's h value:     ",
            round(h,3),"\n"))
  cat(paste("p value (two-tailed): ",
            round((1-pnorm(abs(h)))*2,2),"\n"))
  cat(paste("Reject Null at 95% level? ",
            abs(h)>=1.96,"\n"))
  invisible(NULL)
}

## alternative version if we have the lmtest() library installed
mydw <- function(obj,iterations=15,exact=NULL,tol=1e-10){
  ##if (!is.loaded(symbol.For("pan")))
  
  mf <- model.frame(obj)
  y <- model.response(mf)
  X <- model.matrix(obj)
  dname <- paste(obj$call[2])
   
  n <- nrow(X)
  if (is.null(exact)) 
    exact <- (n < 100)
  k <- ncol(X)
  res <- residuals(obj)
  dw <- sum(diff(res)^2)/sum(res^2)

  if(require(lmtest)){   ## load lmtest library if not loaded
    A <- diag(c(1, rep(2, n - 2), 1))
    A[abs(row(A) - col(A)) == 1] <- -1
    Q1 <- solve(crossprod(X), tol = tol)
    if (exact) {
      MA <- diag(rep(1, n)) - X %*% Q1 %*% t(X)
      MA <- MA %*% A
      ev <- eigen(MA)$values[1:(n - k)]
      if (any(Im(ev) > tol)) 
        warning("imaginary parts of eigenvalues discarded")
      ev <- Re(ev)
      ev <- ev[ev > tol]
      pval <- .Fortran("pan", as.double(c(dw, ev)), as.integer(length(ev)), 
                       as.double(0), as.integer(iterations), x = double(1))$x
      if ((pval > 1) | (pval < 0)) {
        warning("exact p value cannot be computed (not in [0,1]), approximate p value will be used")
        exact <- FALSE
      }
    }
    if (!exact) {
      XAXQ <- t(X) %*% A %*% X %*% Q1
      P <- 2 * (n - 1) - sum(diag(XAXQ))
      Q <- 2 * (3 * n - 4) - 2 * sum(diag(t(X) %*% A %*% A %*% 
                                          X %*% Q1)) + sum(diag(XAXQ %*% XAXQ))
      dmean <- P/(n - k)
      dvar <- 2/((n - k) * (n - k + 2)) * (Q - P * dmean)
      pval <- pnorm(dw, mean = dmean, sd = sqrt(dvar))
    }
    names(dw) <- "DW"
    RVAL <- list(statistic = dw, method = "Durbin-Watson test", 
                 p.value = pval, data.name = dname)
    class(RVAL) <- "htest"
    return(RVAL)
  }
  else{
    cat(paste("Durbin-Watson Statistic:",
              signif(dw,4),
              "\n"))
    return(dw)
  }
}

################################################################
## visualize impacts over time in AR(1) models
## y_t = beta_0 + lambda y_{t-1} + beta_1 x_t + u_t
################################################################
impulse <- function(beta=1,rho){
  halflife <- -log(2)/log(rho)
  if (halflife < 4)
    r <- 0:10
  else
    r <- 0:(2*round(halflife))    ## time-horizon is 0 thru two halflives
  fx <- beta*(rho^r)
  plot(r,fx,
       type="b",    ## plot both points and a line
       xlab="Time",
       ylab="Impact on y")

  ## show halflife 

  abline(v=halflife)
  mtext(side=3,
        at=halflife,
        paste("Halflife of shock at",round(halflife,1)))
  
  out <- cbind(r,fx)    ## return to user
  out
}

